Review — PR #28416: [MLAS] Add CPU DynamicQuantMatMulFp8 contrib op with MLAS FP8 fallback

Big surface area: new contrib op + schema + shape inference + prepack + a fresh public MLAS API (MlasFp8GemmBatch) + scalar fallback + tests + an unrelated MLAS helper refactor (MlasMultiplyOverflowsSizeT). ~2.4k lines. Functionally it looks correct and the test coverage is genuinely good, but I have substantive concerns about the framing of this as an "MLAS FP8" path and about a few API/perf decisions that will be painful to walk back once shipped.

Verdict: needs work before merge. Not because of correctness, but because (a) the MLAS public API being introduced is shipped without any optimized backend behind it, (b) the data-params struct will be ABI-fragile, and (c) several validations get redundantly executed on every Compute().

Substantive concerns

1. MlasFp8GemmBatch is a scalar reference implementation behind a permanent public ABI.
qgemm_fp8.cpp is O(M·N·K) with FP8 decode (Fp8ByteToFloat switch) on the inner element, no register blocking, no SIMD, no packed B. The PR description acknowledges this ("correctness-oriented scalar fallback, not a production performance kernel") — but it lands in mlas.h as MlasFp8GemmBatch, which downstream code will reasonably assume is the MLAS performance path. There is no KleidiAI dispatch added, no x86 dispatch, no dispatch struct at all; today the API has exactly one implementation and that implementation is a triple loop. A user converting Qwen3 to this op on CPU will get unacceptable throughput and have no diagnostic to indicate "this is the reference path." Two options:

• Rename to MlasReferenceFp8GemmBatch (or keep it internal to the contrib op for now) until at least one optimized backend lands.
• Or land a KleidiAI FP8 dispatch in the same MR so the public API has a real reason to exist.

The current framing risks normalizing a slow path under an MLAS name.

2. MLAS_FP8_GEMM_DATA_PARAMS will be ABI-fragile.
14 fields, Fp8Type is per-batch-element (genuinely unusual — why would batch entries vary mode?), BlocksM/K/N are caller-provided and re-derived/validated inside CheckFp8GemmBatchParams, and the strides (ScaleAStrideM, ScaleAStrideK, ScaleBStrideK, ScaleBStrideN) are also caller-provided. Future variants (per-channel B, per-tensor A, group quant) will require new fields, which is breaking for statically linked consumers. Suggest either (a) hide most of these behind a MLAS_FP8_GEMM_BLOCK_LAYOUT sub-struct, or (b) commit to a versioned config pattern now.

Also: Fp8Type per DataParams but Shape is shared — if a caller varies Fp8Type per batch, the kernel implementation cannot precompute mode-specific decode tables. Make it part of MLAS_FP8_GEMM_SHAPE_PARAMS.

3. mlas_fp8_mode naming and provenance.
Lowercase snake_case enum in a public MLAS header breaks convention (rest of MLAS is MLAS_FOO). The comment "FP8 mode is aligned with Arm KleidiAI format/overflow modes. Defined here to keep MLAS FP8 APIs platform-agnostic" then leaks the implementation alignment into the public name. Rename to MLAS_FP8_MODE and drop the KleidiAI reference.

4. Per-Compute() redundant validation.
ValidateZeroPointValuesAreZero walks the full B zero-point tensor every Run() call, even when B (and therefore the zero-point) is a constant initializer. Same for ValidatePositiveFiniteScales(b_scales, …) when scales came from a constant initializer. In a serving scenario these checks burn CPU on every inference. Either:

• Validate constant inputs once at OpKernel construction / PrePack and cache, or
• Drop the zero-point inputs entirely from the schema since the op is symmetric-only and the kernel rejects non-zero values anyway.

The latter is cleaner: an op whose contract is "all zero points must encode 0.0" probably shouldn't accept zero-point inputs at all. Make symmetric explicit in the schema and remove the input slots.

5. block_size_m attribute that must be 1.
The schema exposes block_size_m, the kernel rejects everything except 1, and shape inference enforces 1. This is a future-compatibility footgun: once a model is in the wild with block_size_m=1 as a baked attribute, you cannot change the default semantics. Either drop the attribute now (the op is row-wise-A by contract) or actually implement block_size_m > 1.

6. Per-Compute() allocations from the temp allocator.
Every Run() allocates: A FP8 buffer (M*K*batch), A scale buffer, optional Y float scratch (when output is fp16/bf16), optional B-scale float-conversion buffer. For a 70B-class model with thousands of MatMul calls per token, this is meaningful allocator churn. At minimum, the A FP8 + A scale pair should be sized once and reused, or you should rely on the arena allocator — but note the bug report I just looked at (#28654) shows the arena being disabled in practice by users, so don't over-rely on it.

7. kPackedBMetadataHasFp8ModeIndex is dead weight.
PrePack always writes 1 to this slot; RestorePackedBMetadata rejects anything else. Drop the field and bump kPackedBMetadataVersion if you want to preserve forward compatibility.

8. b_scales_ ownership gymnastics.
The IAllocatorUniquePtr<float> ↔ IAllocatorUniquePtr<void> round-trip via lambda-wrapped deleters in both PrePack and UseSharedPrePackedBuffers is fragile (get_deleter() is moved-from before release() and the assumption that the deleter stays valid is non-obvious). Either store B scales as IAllocatorUniquePtr<void> + element count throughout, or add a small static_cast helper that does this conversion in one place with a comment about why it's safe.

9. Test Fp8Gemm, ZeroColumnReturnsBeforeWorkItemOverflow doesn't actually test what it claims.
BatchN=2, M=SIZE_MAX, N=0 — the BatchN==0 || M==0 || N==0 early return triggers on N==0 before CheckedWorkItems runs, so the overflow path is never exercised. If you want to exercise the overflow guard, you need N != 0, but then ORT_ENFORCE aborts the process. Either:

• Convert CheckedWorkItems and CheckFp8GemmBatchParams to return Status (preferred — ORT_ENFORCE in MLAS code that's reachable from user-controlled shapes is a DoS vector), or
• Delete the misleading test and rely on the contrib-op-level guards.

10. MlasMultiplyOverflowsSizeT refactor is unrelated.
The migration from kleidiai/mlasi_kleidiai.h to mlasi.h is a sensible cleanup but belongs in its own MR. Mixing it in here makes the diff harder to review and the revert story worse if something regresses. Split.

Nits

• core/common/fp8_common.h carries only Arm copyright but is a generic core utility. Consider Microsoft co-attribution per the repo's usual practice for core/ files. (Doesn't block merge.)
• Fp8ByteToFloat's default: return 0.0f; swallows invalid modes silently. Prefer ORT_THROW or __builtin_unreachable() paired with a debug assert — it's only called after IsValidFp8Mode checks anyway, but the silent zero is the wrong fallback.
• QuantizeBlockwiseFp8ABlockDynamic and QuantizeBlockwiseFp8WithScales each iterate src[row_offset + k] twice (once for max_abs, once for quantize). For bf16/fp16 inputs that means two casts per element. Minor, but the comment says "match KleidiAI" — the layout matches, the perf does not.
• gemm_data.ScaleAStrideK = 1; gemm_data.ScaleAStrideM = blocks_k; and the symmetric B strides are caller-derived and re-validated in MLAS — pick one source of truth.
• Schema lists Inputs (2 - 6) but documents 6 inputs. The optional-input chaining (AddOptionalInputEdge between B_zero_point and Y_scale) is correct but easy for model authors to get wrong. A short example in the schema doc would help.
• Y_zero_point provided without Y_scale is silently accepted. Either enforce both-or-neither at the schema level or document explicitly that Y_zero_point requires Y_scale.
• Tests use 1e-5f absolute tolerance for some cases and 0.5f for others — the 0.5f is reasonable for FP8 quantization noise but is an enormous tolerance; consider asserting a relative error bound instead so silent precision regressions surface.

Bottom line

Land the contrib op, but split the work:

1. This MR: contrib op + schema + tests + an internal (un-exported) scalar FP8 GEMM helper used only by this kernel.
2. Follow-up MR: promote MlasFp8GemmBatch to mlas.h together with at least one optimized backend (KleidiAI SME path is the obvious candidate). At that point lock down the DATA_PARAMS struct shape.
3. Separate MR: the MlasMultiplyOverflowsSizeT refactor.

If splitting is not on the table, then before merge please at minimum: (a) rename the public symbol to make the reference nature explicit, (b) move per-Run() constant-input validation into PrePack, (c) drop the dead kPackedBMetadataHasFp8ModeIndex slot, (d) drop block_size_m attribute or implement it, (e) convert MLAS-side ORT_ENFORCE to Status, (f) rename mlas_fp8_mode → MLAS_FP8_MODE and remove the KleidiAI reference from the public header.